""" Base class for llm model """
from typing import Annotated, List, Optional, Union, Iterator
from abc import ABC, abstractmethod
from llmapi.base.types import *
import enum

class ModelType(enum.Enum):
    CHAT = "chat"
    EMBEDDING = "embedding"

class ChatModel(ABC):
    """ Base class for completion model """
    @property
    def type(self) -> ModelType:
        return ModelType.CHAT
    @property
    @abstractmethod
    def name(self) -> str:
        """ model name """
    @abstractmethod
    def completion(self, 
                   prompt: str,
                   max_tokens: int = 16,
                   temperature: float = 0.7,
                   top_p: float = 1.0,
                   top_k: int = 0,
                   presence_penalty: float = 0.0,
                   frequency_penalty: float = 0.0,
                   stop: Optional[List[str]] = [],
                   echo: bool = False,
                   logprobs: int = 0,
                   stream: bool = False,
                   verbose: Optional[bool] = False,
                   **kwargs) -> Union[Completion, Iterator[CompletionChunk]]:
        """ create completion """

    @abstractmethod
    def chat_completion(self, 
                        messages: List[ChatCompletionMessage],
                        temperature: float = 0.2,
                        top_p: float = 0.95,
                        top_k: int = 40,
                        stream: bool = False,
                        stop: Optional[List[str]] = [],
                        max_tokens: int = 256,
                        verbose: Optional[bool] = False,
        **kwargs) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
        """ create chat completion """

class EmbeddingModel(ABC):
    """ Base class for embedding model """
    @property
    def type(self) -> ModelType:
        return ModelType.EMBEDDING
    @property
    @abstractmethod
    def name(self) -> str:
        """ model name """
    @abstractmethod
    def embedding(self, input: str,**kwargs) -> Embedding:
        """ create embedding """


MODELS: Dict[Annotated[str, "model name"],Union[ChatModel, EmbeddingModel]] = {}